// ClassFactory.cpp
// (c) 2008, Charles Lechasseur
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
// 
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
// 
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

#include "ClassFactory.h"

#include "COM/Server.h"
#include "COM/AutoPtr.h"
#include "Memory/ArrayAutoPtr.h"
#include "Registry/Registry.h"
#include <exception>


namespace RRO {
namespace COM
{

// -- Creator class --

    //
    // Constructor.
    //
    Creator::Creator(void)
    {
    }

    //
    // Copy constructor.
    //
    Creator::Creator(const Creator&)
    {
    }

    //
    // Assignment operator.
    //
    Creator& Creator::operator=(const Creator&)
    {
        return *this;
    }

    //
    // Destructor.
    //
    Creator::~Creator(void)
    {
    }

// -- NullCreatorException class --

    //
    // Constructor.
    //
    NullCreatorException::NullCreatorException(void)
        : Exception()
    {
    }

    //
    // Copy constructor.
    //
    NullCreatorException::NullCreatorException(const NullCreatorException& p_Exception)
        : Exception(p_Exception)
    {
    }

    //
    // Assignment operator.
    //
    NullCreatorException& NullCreatorException::operator=(const NullCreatorException& p_Exception)
    {
        if (this != &p_Exception) {
            Exception::operator=(p_Exception);
        }
        return *this;
    }

    //
    // Destructor.
    //
    NullCreatorException::~NullCreatorException(void)
    {
    }

    //
    // Returns exception's description.
    //
    const char* NullCreatorException::what(void) const
    {
        return "RRO::COM::NullCreatorException";
    }

// -- ClassFactory class --

    //
    // Creates a class factory using a given object Creator.
    //
    // @param p_pCreator Creator to use to create objects in CreateInstance.
    //                   Cannot be null; otherwise a NullCreatorException will be thrown.
    // @param p_ClassID Class ID of the object's CoClass.
    // @param p_ProgID ProgID of the object's CoClass.
    // @param p_ThreadingModel Object threading model; defaults to Apartment.
    // @param p_pDescription Description of the object's CoClass; defaults to null.
    //
    ClassFactory::ClassFactory(Creator* p_pCreator, const CLSID& p_ClassID,
                               const std::string& p_ProgID,
                               ThreadingModel p_ThreadingModel /*= THREADING_MODEL_APARTMENT*/,
                               const char* p_pDescription /*= 0*/)
        : m_apCreator(p_pCreator),
          m_ClassID(p_ClassID),
          m_ClassIDAsString(ClassIDToString(p_ClassID)),
          m_ProgID(p_ProgID),
          m_ThreadingModel(p_ThreadingModel),
          m_apDescription(0)
    {
        if (m_apCreator.get() == 0) {
            throw NullCreatorException();
        }

        if (p_pDescription != 0) {
            m_apDescription.reset(new std::string(p_pDescription));
        }
    };

    //
    // Destructor.
    //
    ClassFactory::~ClassFactory(void)
    {
    };

    //
    // Returns the class ID of the COM object's CoClass.
    //
    const CLSID& ClassFactory::ClassID(void) const
    {
        return m_ClassID;
    }

    //
    // Returns the class ID of the COM object's CoClass, as a string.
    //
    const std::string& ClassFactory::ClassIDAsString(void) const
    {
        return m_ClassIDAsString;
    };

    //
    // Return the ProgID of the COM object's CoClass.
    //
    const std::string& ClassFactory::ProgID(void) const
    {
        return m_ProgID;
    }

    //
    // Returns the COM object's CoClass description. Can be null if the CoClass has no description.
    //
    const char* ClassFactory::Description(void) const
    {
        if (m_apDescription.get() != 0) {
            return m_apDescription->c_str();
        } else {
            return 0;
        }
    }

    //
    // Updates the registry by adding or removing the keys pertaining to the COM object
    // handled by this factory.
    //
    // @param p_Register Whether to register or unregister the factory.
    //
    void ClassFactory::UpdateRegistry(bool p_Register)
    {
        // The registry key structure for COM objects implemented in DLLs is as follow:
        //
        // HKEY_CLASSES_ROOT
        // \- CLSID
        //    \- <guid of CoClass>      ->  (Default)      | REG_SZ = <optional description>
        //       \- InProcServer32      ->  (Default)      | REG_SZ = <path to COM server DLL>
        //                              ->  ThreadingModel | REG_SZ = <COM object threading model>
        //       \- ProgID              ->  (Default)      | REG_SZ = <ProgID of CoClass>
        // \- <ProgID of CoClass>       ->  (Default)      | REG_SZ = <optional description>
        //    \- Clsid                  ->  (Default)      | REG_SZ = <guid of CoClass>

        HKEY hRoot = HKEY_CLASSES_ROOT;
        std::string coClassPath = std::string("\\CLSID\\") + ClassIDAsString();
        std::string progIDPath = std::string("\\") + ProgID();

        if (p_Register) {
            std::string inProcServerPath = coClassPath + "\\InProcServer32";
            std::string coClassProgIDPath = coClassPath + "\\ProgID";
            std::string progIDcoClassPath = progIDPath + "\\Clsid";

            std::string threadingModelAsString;
            switch (m_ThreadingModel)
            {
            case THREADING_MODEL_APARTMENT:
                threadingModelAsString = "Apartment";
                break;

            case THREADING_MODEL_FREE:
                threadingModelAsString = "Free";
                break;

            case THREADING_MODEL_BOTH:
                threadingModelAsString = "Both";
                break;

            default:
                break;  // Use an empty string.
            }

            std::string description = (m_apDescription.get() != 0 ? *m_apDescription : "");

            Registry::SetStringValue(hRoot, coClassPath.c_str(), 0, description.c_str());
            Registry::SetStringValue(hRoot, inProcServerPath.c_str(), 0, "");
                // TODO: above, we should store the path to the COM server DLL.
            Registry::SetStringValue(hRoot, inProcServerPath.c_str(), "ThreadingModel",
                threadingModelAsString.c_str());
            Registry::SetStringValue(hRoot, coClassProgIDPath.c_str(), 0, ProgID().c_str());
            Registry::SetStringValue(hRoot, progIDPath.c_str(), 0, description.c_str());
            Registry::SetStringValue(hRoot, progIDcoClassPath.c_str(), 0, ClassIDAsString().c_str());
        } else {
            Registry::DeleteKey(hRoot, coClassPath.c_str());
            Registry::DeleteKey(hRoot, progIDPath.c_str());
        }
    }

    //
    // Creates a new instance of an object handled to this class factory.
    //
    // @param p_pController Controlling object, when requesting aggregation; may be null.
    // @param p_IID ID of the interface requested on the class factory object. Usually IID_IClassFactory.
    // @param p_ppObject Where to store the resulting object reference.
    //
    STDMETHODIMP ClassFactory::CreateInstance(IUnknown* p_pController,
                                              REFIID p_IID, void** p_ppObject)
    {
        HRESULT res = E_UNEXPECTED;

        if (p_ppObject == 0) {
	        res = E_INVALIDARG;
        } else if (p_pController != 0) {
	        // We do not support aggregation for now.
	        res = CLASS_E_NOAGGREGATION;
        } else {
	        // Ask creator to create object for us.
            AutoPtr<IUnknown> object(m_apCreator->Create(), false);
	        try {
                if (!object.IsNull()) {
			        // Ask object for the required interface and return it.
			        res = object->QueryInterface(p_IID, p_ppObject);
		        } else {
			        *p_ppObject = 0;
			        res = E_FAIL;
		        }
	        } catch (...) {
		        // We're not supposed to throw.
	        }
        }

        return res;
    }

    //
    // Attempts to lock or unlock the server of this class factory in memory.
    // This is implemented using CoLockObjectExternal.
    //
    // @param p_Lock Whether to lock or unlock the server.
    //
    STDMETHODIMP ClassFactory::LockServer(BOOL p_Lock)
    {
        // We will simply simulate an extra active object in the COM server.
        if (p_Lock) {
            Server::Instance().AddObject();
        } else {
            Server::Instance().RemoveObject();
        }

        return S_OK;
    }

    //
    // Turns a class ID into its string representation. (Static)
    //
    std::string ClassFactory::ClassIDToString(const CLSID& p_ClassID)
    {
        std::string idAsString;

        // Call StringFromCLSID to do the conversion. This method is documented as being able to 
        // return an error code if it's out of memory, so watch this.
        OLECHAR* pString = 0;
        HRESULT res = ::StringFromCLSID(p_ClassID, &pString);
        if (res != S_OK) {
            // Out of memory.
            throw std::bad_alloc();
        }

        // Save the returned string in a wstring and free it to avoid leaks.
        std::wstring wideString(pString);
        ::CoTaskMemFree(pString);

        // Now we need to convert that string to single-byte.
        int needed = ::WideCharToMultiByte(CP_ACP, 0, wideString.c_str(), -1, 0, 0, 0, 0);
        if (needed > 0) {
            ArrayAutoPtr<char> ansiString(new char[needed]);
            int conversionRes = ::WideCharToMultiByte(CP_ACP, 0, wideString.c_str(), -1,
                                                      ansiString.get(), needed, 0, 0);
            if (conversionRes > 0) {
                // Converted successfully.
                idAsString.assign(ansiString.get());
            } else {
                // Memory problem.
                throw std::bad_alloc();
            }
        } else {
            // Memory problem.
            throw std::bad_alloc();
        }

        return idAsString;
    }

}; // namespace COM
}; // namespace RRO
